# strategy/strategy.py
from abc import ABC, abstractmethod
from typing import Optional, List, Dict, Any, Tuple
from datetime import datetime, timezone, timedelta

import pandas as pd  # 需要pd来处理IndicatorCalculator返回的时间戳

# 修改导入路径以匹配新的 utils 文件夹结构
from utils.indicators import IndicatorCalculator  # 导入新的指标计算器
from entity.k_lines import KLines, KLine  # 假设 klines_loader 在顶层或正确路径
from entity.position import Position, PositionDirection
from entity.account import Account
from entity.strategy_signal import StrategySignal, SignalType

beijing_tz = timezone(timedelta(hours=8))


class Strategy(ABC):
    def __init__(self, strategy_name: str, params: Optional[dict] = None):
        self.strategy_name = strategy_name
        self.params = params if params is not None else {}
        # 策略在初始化时，不再直接接收 KLines 对象字典，
        # 而是应该接收 IndicatorCalculator 对象的字典，这些对象由Operator创建和管理
        print(f"策略 '{self.strategy_name}' 已初始化，参数: {self.params}")

    def _convert_utc_ms_to_beijing_datetime(self, timestamp_ms: int) -> datetime:
        utc_dt = datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc)
        return utc_dt.astimezone(beijing_tz)

    def _convert_pd_timestamp_to_beijing_datetime(self, pd_ts: pd.Timestamp) -> datetime:
        # pandas Timestamp 默认可能是 naive UTC，或者已经是 aware UTC
        if pd_ts.tzinfo is None:
            dt_utc = pd_ts.to_pydatetime().replace(tzinfo=timezone.utc)
        else:
            dt_utc = pd_ts.to_pydatetime()  # 已经是 aware
        return dt_utc.astimezone(beijing_tz)

    @abstractmethod
    def generate_signal(
            self,
            # 修改点: 接收 IndicatorCalculator 对象的字典
            indicator_calculators: Dict[str, IndicatorCalculator],
            current_execution_kline: KLine,  # 1分钟K线
            execution_klines_obj: KLines,  # 1分钟KLines对象
            current_position: Optional[Position],
            account: Account
    ) -> Optional[StrategySignal]:
        pass

    def get_name(self) -> str:
        return self.strategy_name

    def __str__(self) -> str:
        return f"策略(名称: {self.strategy_name}, 参数: {self.params})"


class MultiTimeframeStrategy(Strategy):  # 类名修改
    def __init__(
            self,
            ema_period_5m: int = 20, rsi_period_5m: int = 8,
            boll_period_5m: int = 20, boll_nbdev_5m: int = 2,
            take_profit_pcent: float = 0.046, stop_loss_pcent: float = 0.004,
            rsi_upper_threshold: float = 90.0, rsi_lower_threshold: float = 84.0,
            max_raise_rate: float = 0.01
    ):
        super().__init__(
            strategy_name="MultiTimeframeStrategy_PDTA",  # 修改名称
            params={  # 参数保持不变
                "ema_period_5m": ema_period_5m, "rsi_period_5m": rsi_period_5m,
                "boll_period_5m": boll_period_5m, "boll_nbdev_5m": boll_nbdev_5m,
                "take_profit_pcent": take_profit_pcent, "stop_loss_pcent": stop_loss_pcent,
                "rsi_upper_threshold": rsi_upper_threshold, "rsi_lower_threshold": rsi_lower_threshold,
                "max_raise_rate": max_raise_rate
            }
        )
        self.ema_period_5m = ema_period_5m
        self.rsi_period_5m = rsi_period_5m
        self.boll_period_5m = boll_period_5m
        self.boll_nbdev_5m = float(boll_nbdev_5m)
        self.take_profit_pcent = take_profit_pcent
        self.stop_loss_pcent = stop_loss_pcent
        self.rsi_upper_threshold = rsi_upper_threshold
        self.rsi_lower_threshold = rsi_lower_threshold
        self.max_raise_rate = max_raise_rate

    def generate_signal(
            self,
            indicator_calculators: Dict[str, IndicatorCalculator],  # 接收指标计算器字典
            current_execution_kline: KLine,  # 这是1分钟K线
            execution_klines_obj: KLines,  # 这是1分钟的KLines对象
            current_position: Optional[Position],
            account: Account
    ) -> Optional[StrategySignal]:

        symbol_to_trade = execution_klines_obj.coin
        current_1m_kline_close_time_bj = self._convert_utc_ms_to_beijing_datetime(current_execution_kline.close_time)

        if current_position and current_position.is_open:
            return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade, comment="已有持仓")

        is_5m_candle_closing_now = ((current_execution_kline.close_time + 1) % (5 * 60 * 1000) == 0)
        if not is_5m_candle_closing_now:
            return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade,
                                  comment="非5分钟决策点")

        # --- 获取5分钟指标计算器 ---
        indicators_5m = indicator_calculators.get("5m")
        if not indicators_5m or indicators_5m.df.empty:
            return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade,
                                  comment="无5分钟指标计算器或数据")

        # --- 定位用于计算指标的5分钟K线的时间戳 ---
        # 我们需要的是由 current_execution_kline 收盘而“完成”的那根5分钟K线的开盘时间
        # 这个5分钟K线的指标值应该是已经计算好的。
        # 5分钟K线的开盘时间戳 (pandas Timestamp, UTC)
        # 5分钟周期的毫秒数
        five_min_interval_ms = 5 * 60 * 1000
        # 计算当前1m K线所属的5m K线的理论开盘时间（作为UTC毫秒时间戳）
        completed_5m_kline_open_time_ms = (
                                                      current_execution_kline.open_time // five_min_interval_ms) * five_min_interval_ms
        # 转换为 pandas Timestamp (UTC) 以便从 IndicatorCalculator 获取指标
        completed_5m_kline_open_ts_pd_utc = pd.to_datetime(completed_5m_kline_open_time_ms, unit='ms', utc=True)

        # 获取前一根5分钟K线的开盘时间戳
        prev_5m_kline_open_time_ms = completed_5m_kline_open_time_ms - five_min_interval_ms
        prev_5m_kline_open_ts_pd_utc = pd.to_datetime(prev_5m_kline_open_time_ms, unit='ms', utc=True)

        # --- 从 IndicatorCalculator 获取指标值 ---
        # K线的指标
        # 注意：IndicatorCalculator.get_indicator_value 的 timestamp 参数需要是 pandas Timestamp
        ema_col_name = f"EMA_{self.ema_period_5m}"
        rsi_col_name = f"RSI_{self.rsi_period_5m}"
        boll_lower_col_name = f"BBL_{self.boll_period_5m}_{self.boll_nbdev_5m}"  # BBL_20_2.0
        boll_upper_col_name = f"BBU_{self.boll_period_5m}_{self.boll_nbdev_5m}"  # BBU_20_2.0

        # 获取"当前完成的5分钟K线" (k) 的指标
        # 我们需要 kline_data_k['close'] 作为 close_5_k
        kline_data_k = indicators_5m.get_kline_and_indicators_at_timestamp(completed_5m_kline_open_ts_pd_utc)
        if not kline_data_k: return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade,
                                                   comment="无法获取5m(k)的K线和指标数据")

        close_5_k = kline_data_k.get('close')
        ema_5_k = kline_data_k.get(ema_col_name)
        rsi_5_8_k = kline_data_k.get(rsi_col_name)
        boll_lower_k = kline_data_k.get(boll_lower_col_name)
        boll_upper_k = kline_data_k.get(boll_upper_col_name)

        # 获取"前一根5分钟K线" (k-1) 的指标
        kline_data_k_minus_1 = indicators_5m.get_kline_and_indicators_at_timestamp(prev_5m_kline_open_ts_pd_utc)
        if not kline_data_k_minus_1: return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj,
                                                           symbol_to_trade, comment="无法获取5m(k-1)的K线和指标数据")

        open_5_k_minus_1 = kline_data_k_minus_1.get('open')
        close_5_k_minus_1 = kline_data_k_minus_1.get('close')
        boll_lower_k_minus_1 = kline_data_k_minus_1.get(boll_lower_col_name)  # 使用相同的列名结构
        boll_upper_k_minus_1 = kline_data_k_minus_1.get(boll_upper_col_name)

        if close_5_k is None or ema_5_k is None or rsi_5_8_k is None or \
                boll_lower_k is None or boll_upper_k is None or \
                open_5_k_minus_1 is None or close_5_k_minus_1 is None or \
                boll_lower_k_minus_1 is None or boll_upper_k_minus_1 is None:
            return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade,
                                  comment="5分钟指标值存在None")

        # --- 计算策略衍生值 ---
        if (boll_upper_k - boll_lower_k) == 0:
            boll_brand = 0.5
        else:
            boll_brand = (close_5_k - boll_lower_k) / (boll_upper_k - boll_lower_k)

        if (boll_upper_k_minus_1 - boll_lower_k_minus_1) == 0:
            last_boll_brand = 0.5
        else:
            last_boll_brand = (close_5_k_minus_1 - boll_lower_k_minus_1) / (boll_upper_k_minus_1 - boll_lower_k_minus_1)

        if open_5_k_minus_1 == 0:
            raise_rate = 0.0
        else:
            raise_rate = abs((close_5_k_minus_1 - open_5_k_minus_1) / open_5_k_minus_1)

        # --- 开仓条件判断 (只做多) ---
        condition1 = close_5_k > ema_5_k
        condition2 = self.rsi_lower_threshold < rsi_5_8_k < self.rsi_upper_threshold
        condition3 = boll_brand < last_boll_brand
        condition4 = raise_rate < self.max_raise_rate

        debug_info = (f"C5k={close_5_k:.2f},EMA5k={ema_5_k:.2f}(C>E:{condition1}), "
                      f"RSI5_8k={rsi_5_8_k:.2f}({self.rsi_lower_threshold:.1f}<RSI<{self.rsi_upper_threshold:.1f}:{condition2}), "
                      f"BB={boll_brand:.3f},LastBB={last_boll_brand:.3f}(BB<LBB:{condition3}), "
                      f"RR={raise_rate:.4f}(<{self.max_raise_rate}:{condition4})")

        if condition1 and condition2 and condition3 and condition4:
            open_price = current_execution_kline.close  # 使用当前1分钟K线的收盘价开仓
            print(
                f"策略 {self.strategy_name} 在 {current_1m_kline_close_time_bj.strftime('%Y-%m-%d %H:%M:%S %Z%z')}: 开仓条件满足. 执行价(1m收盘): {open_price:.4f}. {debug_info}")

            # 订单大小计算
            order_size = 0.01  # 默认最小
            if open_price > 0 and self.stop_loss_pcent > 0:
                risk_per_trade_abs = account.current_balance * 0.01
                potential_loss_per_unit = open_price * self.stop_loss_pcent
                if potential_loss_per_unit > 0:
                    order_size = risk_per_trade_abs / potential_loss_per_unit
            if order_size <= 0: order_size = 0.01  # 确保有值

            take_profit_target = open_price * (1 + self.take_profit_pcent)
            stop_loss_target = open_price * (1 - self.stop_loss_pcent)

            return StrategySignal(
                signal_type=SignalType.OPEN_LONG, timestamp=current_1m_kline_close_time_bj,
                symbol=symbol_to_trade, price=open_price, size=order_size,
                take_profit_price=take_profit_target, stop_loss_price=stop_loss_target,
                comment=f"5分钟指标开多: TP={take_profit_target:.4f}, SL={stop_loss_target:.4f}. {debug_info}"
            )

        return StrategySignal(SignalType.HOLD, current_1m_kline_close_time_bj, symbol_to_trade,
                              comment=f"5分钟条件未满足. {debug_info}")